#include <src/pcps.h>
#include <src/rdm.h>
#include <src/blas_lapack.h>

//-------------------------------------------------------------------------------
// pcps::compute_hf_rdm_element -- compute the many body hartree-fock RDM 
//                                 element corresponding to the specified pairs
//                                 of creation and destruction operators
//-------------------------------------------------------------------------------

template <class SCALAR>
SCALAR pcps::compute_hf_rdm_element(const int nops,
                                    const int * const occs,
                                    const SCALAR * const rdm_data,
                                    SCALAR * const det,
                                    int * const ipiv)
{

  // build the determinant matrix
  for (int i = 0; i < nops; i++)
    // hole row
    if ( occs[i] == 0 )
      for (int j = 0; j < nops; j++)
        det[i*nops+j] = (i == j ? 1.0 : 0.0) * pcps::unity<SCALAR>() - rdm_data[i*nops+j];
    // particle row
    else
      for (int j = 0; j < nops; j++)
        det[i*nops+j] = rdm_data[i*nops+j];

  // compute the LU decomposition of the determinant matrix
  int info = 0;
  pcps::xgetrf(nops, nops, det, nops, ipiv, info);
  if (info < 0)
    throw pcps::Exception( (boost::format("pcps::xgetrf failed with error code %i in pcps::compute_hf_rdm_element") % info).str() );

  // compute the determinant value from the LU decomposition's diagonal elements and the signs created by row pivoting
  SCALAR retval = pcps::unity<SCALAR>();
  for (int i = 0; i < nops; i++) {
    retval *= det[i*nops+i];
    if ( i+1 != ipiv[i] )
      retval = -retval;
  }

  // return the value of the determinant
  return retval;

}

// instantiate the template
template double pcps::compute_hf_rdm_element(const int nops,
                                             const int * const occs,
                                             const double * const rdm_data,
                                             double * const det,
                                             int * const ipiv);
template std::complex<double> pcps::compute_hf_rdm_element(const int nops,
                                                           const int * const occs,
                                                           const std::complex<double> * const rdm_data,
                                                           std::complex<double> * const det,
                                                           int * const ipiv);

class n_integer_solutions_key {
  private:
    int _n, _r, _k;
  public:
    n_integer_solutions_key(const int n, const int r, const int k) : _n(n), _r(r), _k(k) {}
  public:
    static bool comp(const n_integer_solutions_key & v, const n_integer_solutions_key & w) {
      if (v._n < w._n) return true;
      if (v._r < w._r) return true;
      if (v._k < w._k) return true;
      return false;
    }
};

//-------------------------------------------------------------------------------
// pcps::n_integer_solutions -- Computes the number of solutions to the equation
//                              x(1) + x(2) + ... + x(n) = r
//                              when the variables x(i) are constrained to be
//                              integers in the range (0, 1, 2, ..., k)
//-------------------------------------------------------------------------------

double pcps::n_integer_solutions(const int n, const int r, const int k) {

  // map to remember past results
  bool(*compare_ptr)(const n_integer_solutions_key &, const n_integer_solutions_key &) = n_integer_solutions_key::comp;
  static std::map<n_integer_solutions_key, double, bool(*)(const n_integer_solutions_key &, const n_integer_solutions_key &)> past_results(compare_ptr);

  // if the result has already been computed, return it
  n_integer_solutions_key key(n,r,k);
  {
    std::map<n_integer_solutions_key, double, bool(*)(const n_integer_solutions_key &, const n_integer_solutions_key &)>::const_iterator it;
    it = past_results.find(key);
    if (it != past_results.end())
      return it->second;
  }

  // prepare a vector to hold the number of variables having each allowed value
  static std::vector<int> n_with_value;
  if (n_with_value.size() < k+1)
    n_with_value.resize(k+1);
  for (int i = 0; i <= k; i++)
    n_with_value[k] = 0;

  // initialize the return value
  double retval = 0.0;

  // Loop over all possible distributions of variables among the values.
  // Note that we do not directly track of how many variables are equal to zero,
  // as this is known by how many variables take on other values.
  while (true) {

    // compute the number of nonzero variables
    int n_nonzero = 0;
    for (int i = 1; i <= k; i++)
      n_nonzero += n_with_value[i];

    // compute the sum of the variables
    int sum = 0;
    for (int i = 1; i <= k; i++)
      sum += i * n_with_value[i];

    // if this distribution solves the equation, count how many ways it can occur
    if (sum == r && n_nonzero <= n) {

      // determine how many variables are nonzero
      int t = 0;
      for (int i = 1; i <= k; i++)
        t += n_with_value[i];

      // count how many ways the variables can satisfy this distribution
      double occurances = pcps::n_choose_m(n, t);
      for (int i = 1; i < k; i++) {
        occurances *= pcps::n_choose_m(t, n_with_value[i]);
        t -= n_with_value[i]; // t is now equal to the number of variables greater than i
      }

      // record how many ways the variables satisfy this distribution
      retval += occurances;

    }

    // increment to the next distribution of variables
    int p;
    for (p = k; p > 0; p--)
      if (++n_with_value[p] > n)
        n_with_value[p] = 0;
      else
        break;

    // stop iterating if all distributions have been processed
    if (p == 0) break;

  }

  // remember the result
  past_results[key] = retval;

  // return the result
  return retval;

}
